import tensorflow as tf
from module import conv2d, maxpool2d, fc


# LeNet，32*32*3
def lenet(input_tensor, num_classes, regularizer=None):
    """
    :param input_tensor: 输入的四维张量，形如[batch_size, height, width, channels]，一般为[batch_size, 32, 32, 3]
    :param num_classes: 分类数
    :param regularizer: 正则化项
    :return: 未经过softmax的一维向量，统计后的模型参数量
    """
    # 第一层，卷积层
    conv1 = conv2d(input_tensor, 'conv1', output_channel=20, kh=5, kw=5, dh=1, dw=1, padding='VALID')
    pool1 = maxpool2d(conv1, 'pool1', kh=2, kw=2, dh=2, dw=2, padding='VALID')
    # 第二层，卷积层
    conv2 = conv2d(pool1, 'conv2', output_channel=50, kh=5, kw=5, dh=1, dw=1, padding='VALID')
    pool2 = maxpool2d(conv2, 'pool2', kh=2, kw=2, dh=2, dw=2, padding='VALID')
    # 展开张量
    shape = pool2.get_shape()
    nodes = shape[1].value * shape[2].value * shape[3].value
    reshaped = tf.reshape(pool2, [-1, nodes], name='reshape')
    # 第三层，全连接层
    fc3 = fc(reshaped, 'fc3', 120)
    # 第四层，全连接层
    fc4 = fc(fc3, 'fc4', 84)
    # 第五层，全连接映射层，保证softmax的输入输出维度一致
    logit = fc(fc4, 'out', num_classes, activation=False, regularizer=regularizer)
    return logit
